#ifndef _KMEANS_H_
#define _KMEANS_H_

#include <vector>
#include "Index.h"
#include "Cluster.h"
#include "DocumentVector.h"

class Kmeans 
{

private:

	int num_clusters;
	std::vector<Cluster> clusters;
	std::vector<DocumentVector> documents;
	Index index;

public:

	Kmeans(Index index_, std::vector<DocumentVector> documents_, int k) : num_clusters(k), documents(documents_), index(index_) {}	
	
	std::vector<Cluster> run() 
	{
		// initialization
		for(int i = 0; i < num_clusters; i++)
			clusters.insert(clusters.end(), Cluster(documents[i], index));
		
		//stop the algorithm after 5 iterations
		for (int a = 0; a < 5; a++) 
		{
			for(int i = 0; i < num_clusters; i++) 
				clusters[i] = Cluster(clusters[i].get_centroid(), index);
		
			// put each document in the appropriate cluster
			for(unsigned int i = 0; i < documents.size(); i++) 
			{
				int max_index = 0;
				double max_similarity = 0;
			
				for(int j = 0; j < num_clusters; j++) 
				{
					if (clusters[j].similarity(documents[i]) > max_similarity)
					{
						max_index = j;
						max_similarity = clusters[j].similarity(documents[i]);
					}	
				}
				// cout << "max index: " << max_index << endl;
				
			 	clusters[max_index].add_document(documents[i]);
			}
					
			for(int i = 0; i < num_clusters; i++)
				clusters[i].update_centroid();
			
		}
		
		return clusters;
	}
};

#endif /* _KMEANS_H_ */
